﻿#include "recursion.h"
#include "storage.h"
#include "dns_base.h"
table_t *rnode_table = NULL;

static int add_to_record(void ** record, unsigned char * key, void * val)
{
	rnode_event* _record = (rnode_event*)*record;
	request_node* _val = (request_node*)val;

	if (!_record)
	{
		_record = (rnode_event*)malloc(sizeof(rnode_event));
		_record->nodes = malloc(32 * sizeof(request_node*));
		_record->nodes[0] = _val;
		_record->len = 1;
		_record->cap = 32;
		pthread_mutex_init(&_record->_lock, NULL);
		*record = _record;
	}
	else
	{
		pthread_mutex_lock(&_record->_lock);
		_record->nodes[_record->len] = _val;
		++_record->len;
		if (_record->len >= _record->cap)
		{

			_record->cap += 32;
			_record->nodes = realloc(_record->nodes, _record->cap * sizeof(request_node*));
		}
		pthread_mutex_unlock(&_record->_lock);
	}

	return _record->len;
}


inline table_t* get_rnodes()
{
	if (!rnode_table)
	{
		rnode_table = create_table(1024);
		rnode_table->add_value = add_to_record;
	}
	return rnode_table;
}



request_node* generate_request_node(request_info* pinfo)
{
	if (pinfo == NULL)
		return NULL;
	request_node *prn = malloc(sizeof(request_node));
	init_request_node(prn);
	prn->next = NULL;
	pthread_mutex_lock(&pinfo->_lock);
	request_node **tail = &(pinfo->node);
	while (*tail != NULL)
	{
		//prn->prev = *tail;
		tail = &((*tail)->next);
	}

	*tail = prn;
	pthread_mutex_unlock(&pinfo->_lock);
	prn->parent = pinfo;
	return prn;
}


void remove_update_request(request_node* pnode, rnode_event *rnodes)
{
	request_info* pinfo = pnode->parent;
	pthread_mutex_lock(&pinfo->_lock);
	request_node **tail = &(pinfo->node);
	int itmp = 0;
	while (*tail != NULL)
	{
		itmp = 0;
		rnode_event* _record = (rnode_event*)find_from_table(get_rnodes(), (*tail)->lowerdomain.label[0], (*tail)->lowerdomain.label_len[0], 0);
		if (_record)
		{
			pthread_mutex_lock(&_record->_lock);
			int i;
			for (i = 0; i < _record->len; ++i)
			{
				if (_record->nodes[i] == (*tail))
				{
					_record->nodes[i] = NULL;
					//break;
				}
			}
			pthread_mutex_unlock(&_record->_lock);
		}

		for (itmp = 0; itmp < rnodes->len; ++itmp)
		{
			if (rnodes->nodes[itmp] == (*tail))
			{
				rnodes->nodes[itmp] = NULL;
				//break;
			}
		}


		*tail = (*tail)->next;


	}
	pthread_mutex_unlock(&pinfo->_lock);
	free_request(pinfo);
}

void free_request(request_info* pinfo)
{
	if (pinfo == NULL)
		return;

	pthread_mutex_lock(&pinfo->_lock);
	request_node **tail = &(pinfo->node);
	while (*tail != NULL)
	{
		request_node *tmpr = *tail;
		*tail = (*tail)->next;
		free(tmpr);

	}
	pthread_mutex_unlock(&pinfo->_lock);
	free(pinfo);
}


int resp_to_client(request_node* rnode, unsigned char* rname, mvalue_t *mv, unsigned char *sndbuf, int sndlen)
{
	if (!rnode || !rnode->parent || rnode->parent->caddr.sin_port == 0)
	{
		return 0;
	}

	dns_info_t* dt = create_dns_info(sndbuf, sndlen);

	dt->header->ra = 1;
	dt->header->qr = 1;
	dt->header->rd = 1;
	if (mv)
	{
		int i = 0;
		request_node *tmpnode = rnode;

		dns_fill_query(dt, tmpnode->parent->origin_domain, A, CLASS_IN);

		while (tmpnode && tmpnode->prev_chain && tmpnode->prev_chain->hascname == 1)
		{
			dns_fill_an(dt, tmpnode->prev_chain->lowerdomain.domain, CNAME, IN, 86400, strlen(tmpnode->lowerdomain.domain) + 1, tmpnode->lowerdomain.domain);
			tmpnode = tmpnode->prev_chain;
		}

		dt->header->id = tmpnode->id;

		//if (rnode->hascname == 1)
		//{
			//dns_fill_an(dt, rnode->lowerdomain.label[0], CNAME, IN, 86400, strlen(rname)+1, rname);
		//}
		unsigned char *buff = mv->buff;
		for (; i < mv->num; i++)
		{
			// 			struct sockaddr_in aaddr;
			// 			make_addr_from_bin(&aaddr, buff + (i * 4));
			// 			char* tn = inet_ntoa(aaddr.sin_addr);
			// 			printf("resp_to_client  addr %s\n", tn);
			dns_fill_an(dt, rname, A, IN, 86400, 4, buff + (i * 4));

		}
		dns_fill_header(dt);
		int res = sendto(rnode->parent->cfd, dt->buff, dt->start, 0, (struct sockaddr*)&rnode->parent->caddr, sizeof(struct sockaddr));
		dns_fill_header(dt);
		free_dns_info(dt);
		return res;
	}
}



// 返回
// 结速处理
// 查找ns 
int is_ended(dns_info_t*info)
{
	dns_record_t* record = NULL;
	unsigned char * name = info->querys[0].qname;
	if (info->header->ans_count > 0)
	{
		int i = 0, j = 0;
		for (i = 0; i < info->header->ans_count; i++)
		{
			for (j = 0; j < info->header->ans_count; j++)
			{
				if (info->records[j].par->type == SOA)  //请求出错 结速处理
				{
					return 0;
				}
				if (strcmp(name, info->records[j].name) == 0)
				{
					if (info->records[j].par->type == CNAME)
					{
						name = info->records[j].rdata;
						unsigned char * buff = (unsigned char *)malloc(info->records[j].name_len + sizeof(dns_record_t) + info->records[j].par->rdata_len);
						dns_record_t*rec = (dns_record_t*)buff;
						rec->par = NULL;
						rec->name = buff + sizeof(dns_record_t);
						rec->rdata = buff + sizeof(dns_record_t) + info->records[j].name_len;
						memcpy(rec->name, info->records[j].name, info->records[j].name_len);
						memcpy(rec->rdata, info->records[j].rdata, info->records[j].par->rdata_len);
					}
						
					break;
				}
			}
			dns_record_t temp = info->records[i];
			info->records[i] = info->records[j];
			info->records[j] = temp;
		}
		record = &info->records[info->header->ans_count - 1]; //取出最后一个

		if (record->par->type == info->querys[0].par->qtype)   //返回错误  结速处理
		{
			return 0;
		}
		if (record->par->type == AAAA && info->querys[0].par->qtype == A)  //Ipv6 的返回型式 也可以结速处理
		{
			return 0;
		}
	}
	return 1;
}


int get_ns_a_from_dns_info(dns_info_t*info,unsigned char * ip, unsigned int ip_len)
{
	int ip_num = 0;
	if (info->header->auth_count > 0 && info->header->add_count > 0)  //使用自带NS和NS的A 接着查找   
	{
		int i, j;
		dns_record_t * auth_record, *add_record;
		auth_record = &info->records[info->header->ans_count];
		add_record = &info->records[info->header->ans_count + info->header->auth_count];
		for (i = 0; i < info->header->auth_count; i++)
		{
			for (j = 0; j < info->header->add_count; j++)
			{
				if (add_record[j].par->type == A && strcmp(add_record[j].name, auth_record[i].name) == 0)
				{
					if (ip_num * 4 > ip_len)
					{
						break;
					}
					memcpy(ip + ip_num * 4, add_record[j].rdata, 4);
					ip_num++;
				}
			}
		}
	}
	return ip_num;
}





int update_event_ex(int afd, dns_info_t* q, unsigned char *sndbuf, int sndlen)
{
	rnode_event *rnodes = (rnode_event*)take_from_table(get_rnodes(), q->querys->qname, q->querys->qname_len, 0);
	int i;
	if (!rnodes)
		return;



	for (i = 0; i < rnodes->len; ++i)
	{
		if (rnodes->nodes[i])
		{



		}

	}


	unsigned int ip_num, ip_len = 128;
	unsigned char ip[128];
	unsigned char ns[512];

	dns_record_t * auth_record, *add_record;
	auth_record = &q->records[q->header->ans_count];
	add_record = &q->records[q->header->ans_count + q->header->auth_count];
	packet_type temp_packet;


	if (is_ended(q))
	{
		if (q->header->auth_count > 0 && q->header->add_count > 0) //自带NS和ＮＳ的Ａ记录
		{
			ip_num = get_ns_a_from_dns_info(q, ip, ip_len);
		}
		if (q->header->auth_count > 0 && ip_num <= 0) //取ＮＳ查找本地没有就  创建支线请求 ＮＳ的Ａ记录
		{
			for (i = 0; i < q->header->auth_count; i++)
			{
				int num = find_a_ip(auth_record[i].rdata, ip + ip_num * 4, ip_len - ip_num * 4);
				if (num == 0)
				{
					unsigned char* _ns_ip = ip + ip_num * 4;
					unsigned int _ns_ip_len = ip_len - ip_num * 4;
					check_dns_name(auth_record[i].rdata, temp_packet);
					int num = find_ns_ip(&temp_packet, _ns_ip, _ns_ip_len);
					int j = 0;
					for (j = 0; j < num; ++j)
					{
						dns_info_t* info = create_query(auth_record[i].rdata, A, sndbuf, sndlen);
						query_from_auth_server_ex(info, afd, _ns_ip + j * 4, 53);
					}
				}
				else
					ip_num += num;
			}
		}
		if (ip_num <= 0) //  没有带ｎｓ　只能从根查起
		{
			check_dns_name(q->records[q->header->ans_count].rdata, temp_packet);
			ip_num = find_ns_ip(&temp_packet, ip, ip_len);
			int j = 0;
			for (j = 0; j < ip_num; ++j)
			{
				dns_info_t* info = create_query(q->records[q->header->ans_count - 1].rdata, q->querys[0].par->qtype, sndbuf, sndlen);
				query_from_auth_server_ex(info, afd, ip + j * 4, 53);
			}
		}

	}

}


request_node*  add_rnode(request_info* rinfo,packet_type*lowerdomain)
{
	request_node* rnode = generate_request_node(rinfo);
	cpy_dns_name(&rnode->lowerdomain, lowerdomain);
	insert_to_table(get_rnodes(), rnode->lowerdomain.label[0], rnode->lowerdomain.label_len[0], (void*)rnode, 0);
	return rnode;
}



void add_to_ns_list()
{

}




int find_a_ip(unsigned char * key, unsigned char*ip, unsigned int ip_len)
{
	unsigned int ip_num = 0;
	mvalue_t * mv  = find_value_from_dataset(key, A);
	if (mv)
	{
		if (mv->num *4 > ip_len)
		{
			return 0;
		}
		memcpy(ip, mv->buff, mv->num * 4);
		ip_num = mv->num;
	}

	return ip_num;
}


int  find_ns_ip(packet_type*lowerdomain, unsigned char * ip, unsigned int ip_len)
{	
	
	unsigned int label_count = 0;
	int ip_num = 0;
	mvalue_t* mv = NULL;
	unsigned char * itor = lowerdomain->label[label_count];
	while (1)
	{
		mv = find_value_from_dataset(itor, NS);
		if (mv != NULL)
			break;

		label_count++;
		if (label_count > lowerdomain->label_count)
			return -1;

		if (label_count == lowerdomain->label_count) // root
		{
			itor = lowerdomain->label[label_count - 1] + lowerdomain->label_len[label_count - 1];
			itor[1] = '\0';
		}
		else
		{
			itor = lowerdomain->label[label_count];
		}
	}
	if (mv)
	{
		int i = 0;
		mvalue_t* val_A;
		unsigned char *buff = mv->buff;
		for (i = 0; i < mv->num; i++)
		{
			val_A = find_value_from_dataset(buff, A);
			buff += strlen(buff) + 1;
			if (val_A)
			{
				int num = val_A->num;

				if (val_A->num > ip_len - ip_num)
				{
					num = ip_len;
				}
				memcpy(ip + (ip_num * 4), val_A->buff, num * 4);
				ip_num += num;
			}
			if (ip_num * 4 >= ip_len) //达到最大的ip 数量
			{
				break;
			}

		}

	}
	return ip_num;
}



void update(int afd, dns_info_t* q, request_node* rnode, unsigned char *sndbuf, int sndlen, rnode_event *rnodes)
{
	memset(sndbuf, 0, 512);
	int ret = 0;
	int label_count = 0;
	packet_type *lowerdomain = &rnode->lowerdomain;
	unsigned char *itor = lowerdomain->label[label_count];
	int td_len = lowerdomain->label_len[label_count];
	mvalue_t *mv;
	unsigned char *ip[320] = { 0, };
	//int hash = &(lowerdomain->hash[label_count]);
	int tmpns = 0;
	// TODO
	if (q->header->ans_count > 0 && q->records[0].par->type == A)
	{
		// ???
		int ip_num = 0, pos = 0, i = 0;
		unsigned char*key = NULL;

		for (; i < q->header->ans_count; i++)
		{
			if (key == NULL)
			{
				key = q->records[i].name;
			}
			else if (strcmp(key, q->records[i].name) != 0)
			{
				mvalue_t* val = create_mvalue(ip_num * sizeof(int));
				unsigned char* buff = val->buff;
				val->type = _A;
				int j = pos;
				for (; j < i; j++)
				{
					*((int*)buff) = *((int*)q->records[j].rdata);
					buff += sizeof(int);
					val->num++;
				}
				pos = i;
				key = q->records[i].name;
			}
			ip_num++;
		}
		mv = create_mvalue(ip_num * sizeof(int));
		unsigned char* buff = mv->buff;
		int j = pos;
		for (; j < i; j++)
		{
			*((int*)buff) = *((int*)q->records[j].rdata);
			buff += sizeof(int);
			mv->num++;
		}
	}
	else
	{
		mv = find_value_from_dataset(itor, A);
	}

	packet_type dst_lowerdomain;
	dst_lowerdomain.label_count = 0;
	if (mv == NULL && q->header->ans_count > 0 && q->records[0].par->type == CNAME)
	{
		// ?????DNS?? A????????????????CNAME????CNAME??A
		// CNAME

		check_dns_name(q->records[0].rdata, &dst_lowerdomain);
		lowerdomain = &dst_lowerdomain;
		label_count = 0;
		td_len = lowerdomain->label_len[label_count];
		itor = q->records[0].rdata;
		rnode->hascname = 1;
		// ?????????????CNAME??A
		mv = find_value_from_dataset(itor, A);
	}
	// ????A,????????
	if (mv)
	{
		// ???????????A??????????CNAME???????????
		// ?????cname itor ????????????????????cname  itor????cname?????????
		if (rnode->prev_chain == NULL)
		{
			// TODO ?????????...

			resp_to_client(rnode, itor, mv, sndbuf, sndlen);
			remove_update_request(rnode, rnodes);
			return;
		}
		else
		{

			// ?и????
			// ????????????CNAME?????????????????????????NS
			request_node *tmpnode = rnode->prev_chain;
			while (tmpnode->hascname == 1 && tmpnode->prev_chain != NULL)
			{
				tmpnode = tmpnode->prev_chain;
			}
			if (tmpnode->prev_chain == NULL)
			{
				// TODO ?????????...
				resp_to_client(rnode, itor, mv, sndbuf, sndlen);
				remove_update_request(rnode, rnodes); //free ????
				return;
			}
			else
			{
				// NS ??CNAME   ????????
				int j;
				//int tmpn = mv->num > 2 ? 2 : mv->num;
				for (j = 0; j < mv->num; ++j)
				{
					struct sockaddr_in aaddr;
					make_addr_from_bin(&aaddr, mv->buff + j * 4);
					aaddr.sin_port = htons(53);
					tmpnode->afd = afd;
					query_from_auth_server(tmpnode, sndbuf, 512, &aaddr);

				}
				if (j > 0)
				{
					insert_to_table(get_rnodes(), tmpnode->lowerdomain.label[0], tmpnode->lowerdomain.label_len[0], (void*)tmpnode, 0);
				}
				return;
			}
		}
	}
	else
	{
		// A??????????????NS????CNAME
		// ???itor????????cname itor ????????????????????cname  itor????cname?????????
		// ??
		packet_type tmppt;
		tmppt.label_count = 0;
	rec:
		while (1)
		{
			mv = find_value_from_dataset(itor, NS);
			if (mv != NULL)
				break;

			label_count++;
			if (label_count > lowerdomain->label_count)
				return -1;

			if (label_count == lowerdomain->label_count) // root
			{
				itor = lowerdomain->label[label_count - 1] + lowerdomain->label_len[label_count - 1];
				itor[1] = '\0';
				td_len = 2;
				//thash = 0;
				//hash = &thash;
			}
			else
			{
				itor = lowerdomain->label[label_count];
				td_len = lowerdomain->label_len[label_count];
				//hash = &(lowerdomain->hash[label_count]);
			}
		}
		if (mv)
		{
			int ip_num = retrive_ip(mv, ip, 320);
			if (ip_num > 0)
			{
				int j = 0;
				unsigned char *buff = mv->buff;
				int len = 0;
				//int tmpn = ip_num > 2 ? 2 : ip_num;
				request_node* sub_node = NULL;
				// NS??A???????????????cname  ????????????NS??A????cname??A
				if (dst_lowerdomain.label_count > 0)
				{
					sub_node = generate_request_node(rnode->parent);
					sub_node->prev_chain = rnode;
					//sub_node->lowerdomain = *lowerdomain;
					//memcpy(&sub_node->lowerdomain, lowerdomain, sizeof(packet_type));
					cpy_dns_name(&sub_node->lowerdomain, lowerdomain);
					sub_node->afd = afd;
				}

				for (j = 0; j < ip_num; ++j)
				{
					struct sockaddr_in aaddr;
					make_addr_from_bin(&aaddr, ip + j * 4);
					aaddr.sin_port = htons(53);

					// NS?? CNAME ????  NS IP?????  ?????????????
					if (dst_lowerdomain.label_count > 0)
					{
						query_from_auth_server(sub_node, sndbuf, sndlen, &aaddr);
					}
					else
					{
						// NS A ????????????????
						//char* tn = inet_ntoa(aaddr.sin_addr);
						//printf("zzz  addr %s\n", tn);
						query_from_auth_server(rnode, sndbuf, 512, &aaddr);
					}

				}
				if (sub_node)
				{
					insert_to_table(get_rnodes(), sub_node->lowerdomain.label[0], sub_node->lowerdomain.label_len[0], (void*)sub_node, 0);
				}
				else
				{
					insert_to_table(get_rnodes(), rnode->lowerdomain.label[0], rnode->lowerdomain.label_len[0], (void*)rnode, 0);
				}
			}
			else
			{
				// NS IP????? ???°? ????NS


				itor = mv->buff;
				label_count = 0;
				lowerdomain = &tmppt;
				check_dns_name(itor, lowerdomain);
				dst_lowerdomain.label_count = 1;
				goto rec;
			}
		}
	}
}


int update_event(int afd, dns_info_t* q, unsigned char *sndbuf, int sndlen)
{
	rnode_event *rnodes = (rnode_event*)take_from_table(get_rnodes(), q->querys->qname, q->querys->qname_len, 0);
	int i;
	if (!rnodes)
		return;
	for (i = 0; i < rnodes->len; ++i)
	{
		if (rnodes->nodes[i])
		{
			update(afd, q, rnodes->nodes[i], sndbuf, sndlen, rnodes);
		}

	}
	free(rnodes);
}

int query_from_auth_server(request_node *pnode, unsigned char *sndbuf, int sndlen, struct sockaddr_in *aaddr)
{
	memset(sndbuf, 0, sndlen);
	if (pnode == NULL)
		return -1;
	if (pnode->lowerdomain.label_offsets[0] > 256)// prn->cur_domain ??256
		return  -1;
	if (pnode->id == 0)
	{
		pnode->id = abs(get_hash(pnode->lowerdomain.label[0], pnode->lowerdomain.label_len[0]));
	}
	dns_info_t* dt = create_dns_info(sndbuf, sndlen);
	dns_fill_query(dt, pnode->lowerdomain.label[0], A, CLASS_IN);
	dt->header->id = pnode->id;
	dt->header->rd = 1;
	dns_fill_header(dt);

	int res = sendto(pnode->afd, sndbuf, dt->start, 0, (struct sockaddr*)aaddr, sizeof(struct sockaddr_in));
	dns_fill_header(dt);
	free_dns_info(dt);
}




dns_info_t* create_query(unsigned char *qname, unsigned int qtype, unsigned char *sendbuf, unsigned int sendlen)
{
	dns_info_t* dt = create_dns_info(sendbuf, sendlen);
	dns_fill_query(dt, qname, qtype, CLASS_IN);
	dt->header->id = abs(get_hash(qname, strlen(qname) + 1));
	dt->header->rd = 1;
	dns_fill_header(dt);
	return dt;
}


int query_from_auth_server_ex(dns_info_t *info, unsigned int fd, unsigned char * ip,unsigned int port)
{
	struct sockaddr_in aaddr;
	make_addr_from_bin(&aaddr, ip);
	aaddr.sin_port = htons(port);
	int res = sendto(fd, info->buff, info->start, 0, (struct sockaddr*)&aaddr, sizeof(struct sockaddr_in));
	dns_fill_header(info);
	free_dns_info(info);
}

int query_from_auth_server_dnspod(request_info *pinfo, unsigned char *sndbuf, int sndlen, struct sockaddr_in *aaddr)
{
	if (!pinfo)
		return -1;
	memset(sndbuf, 0, sndlen);
	//if (pinfo->lowerdomain.label_offsets[0] > 256)// prn->cur_domain ??256
		//return  -1;
	if (pinfo->id == 0)
	{
		pinfo->id = abs(get_hash(pinfo->qbuffer, strlen(pinfo->qbuffer)));
	}
	dns_info_t* dt = create_dns_info(sndbuf, sndlen);
	dns_fill_query(dt, pinfo->qbuffer, pinfo->qtype, CLASS_IN);
	dt->header->id = pinfo->id;
	dt->header->rd = 1;
	dns_fill_header(dt);

	int res = sendto(pinfo->afd, sndbuf, dt->start, 0, (struct sockaddr*)aaddr, sizeof(struct sockaddr_in));
	dns_fill_header(dt);
	free_dns_info(dt);
}